from typing import Dict, List

from pyspark import Row
from pyspark.ml.param import Params, Param
from pyspark.ml.param.shared import HasInputCols
from pyspark.sql import DataFrame

from gai.v2.spark.base import SparkTransformer
from gai.v2.utils import get_or_create_spark_session

from gai.v2.spark.demo_data import demo_data_for_cell_nullifier


def make_nullify_cells(columns: List[str], null_equiv: dict):
    """Given a list of column names and a dictionary specifying the values that
    are equivalent to null values, return a mapper function that transforms a
    list of rows so that the cells equivalent to null are replaced with null.

    Args:
        columns:
            the columns that are checked for nullity.
        null_equiv:
            the dictionary of null-equivalent specifiers.

    Returns:
        the mapper function.
    """
    false = lambda x: False

    def nullify_cell(rows):
        """Given a list of rows, fill each cell that satisfies ``null_equiv``
        with null and return the modified rows.
        """
        for row in rows:
            row_dict = row.asDict()
            for col in columns:
                if row_dict[col] is not None and null_equiv.get(col, false)(row_dict[col]):
                    row_dict[col] = None
            yield Row(**row_dict)

    return nullify_cell


def nullify_cells(df: DataFrame, null_equiv: Dict, columns: List[str] = None):
    """Given a dataframe, a list of column names and a dictionary specifying
    equivalents to null values, return a dataframe whose cells within the
    specified columns are replaced with null if they are equivalent to null
    values.

    Args:
        df:
            the dataframe.
        null_equiv:
            the dictionary that specifies the values that are equivalent to null
            values. Each (key, value) pair specifies the column name and the
            values that are equivalent to null. The latter can be specified
            via a list (set, tuple, etc.) of values equivalent to null, or a
            unary predicate that returns ``True`` on values equivalent to null
            values. If no value is specified for a column, the default null
            value (namely, ``None``) is assumed.
        columns:
            the list of columns that are checked.

    Returns:
        the result dataframe.

    See Also:
        :class:`CellNullifier`
    """

    null_equiv_copy = null_equiv.copy()

    # transform non-callable specifiers to callable
    for name, equiv in null_equiv_copy.items():
        if not callable(equiv):
            null_equiv_copy[name] = lambda x: x in equiv

    nullified_rdd = df.rdd.mapPartitions(
        make_nullify_cells(columns or df.columns, null_equiv_copy)
    )
    spark = get_or_create_spark_session()
    nullified_df = spark.createDataFrame(nullified_rdd, df.schema)

    return nullified_df


class CellNullifier(SparkTransformer, HasInputCols):
    """Fills the cells that are equivalent to null with null.

    Args:
        nullEquiv:
            the dictionary that specifies the values that are equivalent to null
            values. Each (key, value) pair specifies the column name and the
            values that are equivalent to null. The latter can be specified
            via a list (set, tuple, etc.) of values equivalent to null, or a
            unary predicate that returns ``True`` on values equivalent to null
            values. If no value is specified for a column, the default null
            value (namely, ``None``) is assumed.
        inputCols:
            the list of columns that are checked. If this argument is
            unspecified, every column is checked for nullity.


    >>> df = demo_data_for_cell_nullifier()
    >>> df.show()
    +------+--------+--------------------+-------------+--------------------+---------------+
    |secret|     day|                gid_|_month_offset|         ft_usertags|ft_category_cnt|
    +------+--------+--------------------+-------------+--------------------+---------------+
    |     1|20181114|ANDROID-68b6c301c...|         null|                    |               |
    |     1|20181114|ANDROID-68b6c301c...|           -1|                    |               |
    |     2|20181114|ANDROID-68b6c301c...|            0|                    |               |
    |     2|20181114|ANDROID-68b6c301c...|           -1|                    |               |
    |     3|20181114|                null|            0|                    |               |
    |     3|20181114|                null|           -1|                    |               |
    |     4|20181126|ANDROID-0000db9c0...|            0|                    |              3|
    |     4|20181126|ANDROID-0000db9c0...|           -1|                    |              3|
    |     5|20181126|ANDROID-846049d18...|            0|022000,026400,02d400|             12|
    |     5|20181126|ANDROID-846049d18...|           -1|022000,026400,02d400|             13|
    |     7|20181126|        NON-EXISTENT|            0|                    |               |
    |     7|20181126|        non-existent|           -1|                    |               |
    |     8|20180709|ANDROID-000340aa5...|            0|    02b000,02b100,h0|             14|
    |     8|20180709|ANDROID-000340aa5...|           -1|    02b000,02b100,h0|             15|
    |     9|20180709|ANDROID-00033a1c3...|            0|022000,022500,02d300|              7|
    |     9|20180709|ANDROID-00033a1c3...|           -1|022000,022100,02d300|              9|
    |    11|20180709|                null|            0|                    |               |
    |    11|20180709|                null|           -1|                    |               |
    +------+--------+--------------------+-------------+--------------------+---------------+
    <BLANKLINE>
    >>> nullifier = CellNullifier({'gid_'           : lambda x: x.casefold() == 'non-existent'.casefold(),
    ...                            'ft_category_cnt': ['']},)
    >>> result = nullifier.transform(df)
    >>> result.show()
    +------+--------+--------------------+-------------+--------------------+---------------+
    |secret|     day|                gid_|_month_offset|         ft_usertags|ft_category_cnt|
    +------+--------+--------------------+-------------+--------------------+---------------+
    |     1|20181114|ANDROID-68b6c301c...|         null|                    |           null|
    |     1|20181114|ANDROID-68b6c301c...|           -1|                    |           null|
    |     2|20181114|ANDROID-68b6c301c...|            0|                    |           null|
    |     2|20181114|ANDROID-68b6c301c...|           -1|                    |           null|
    |     3|20181114|                null|            0|                    |           null|
    |     3|20181114|                null|           -1|                    |           null|
    |     4|20181126|ANDROID-0000db9c0...|            0|                    |              3|
    |     4|20181126|ANDROID-0000db9c0...|           -1|                    |              3|
    |     5|20181126|ANDROID-846049d18...|            0|022000,026400,02d400|             12|
    |     5|20181126|ANDROID-846049d18...|           -1|022000,026400,02d400|             13|
    |     7|20181126|                null|            0|                    |           null|
    |     7|20181126|                null|           -1|                    |           null|
    |     8|20180709|ANDROID-000340aa5...|            0|    02b000,02b100,h0|             14|
    |     8|20180709|ANDROID-000340aa5...|           -1|    02b000,02b100,h0|             15|
    |     9|20180709|ANDROID-00033a1c3...|            0|022000,022500,02d300|              7|
    |     9|20180709|ANDROID-00033a1c3...|           -1|022000,022100,02d300|              9|
    |    11|20180709|                null|            0|                    |           null|
    |    11|20180709|                null|           -1|                    |           null|
    +------+--------+--------------------+-------------+--------------------+---------------+
    <BLANKLINE>
    """

    nullEquiv = Param(Params._dummy(), "nullEquiv",
                      "a dictionary specifying which values are equivalent to "
                      "null, which can be either a list (set, tuple, etc.) or "
                      "a predicate")

    def __init__(self, nullEquiv: Dict, inputCols: List[str] = None):
        super(CellNullifier, self).__init__()

        self.setInputCols(inputCols) \
            .setNullEquiv(nullEquiv)

    def setNullEquiv(self, nullEquiv):
        """Sets the specifier describing which values are equivalent to null.

        Args:
            nullEquiv:
                a dictionary specifying which values are equivalent to null.
        Returns:
            self.
        """
        self._paramMap[self.nullEquiv] = nullEquiv
        return self

    def getNullEquiv(self):
        """

        Returns:
            a dictionary specifying which values are equivalent to null.
        """
        return self.getOrDefault(self.nullEquiv)

    def setCheckedCols(self, checkedCols):
        """

        Args:
            checkedCols:
                the list of columns that are checked.
        Returns:
            self.
        """
        self.setInputCols(checkedCols)
        return self

    def getCheckedCols(self):
        """

        Returns:
            the list of columns that are checked.
        """
        return self.getInputCols()

    def _transform(self, dataset):
        return nullify_cells(dataset, self.getNullEquiv(), self.getCheckedCols() or dataset.columns)
